package com.ls.knowledge.llm.memory;

import cn.hutool.core.util.StrUtil;
import com.ls.knowledge.llm.entity.ChatConversationMessage;
import com.ls.knowledge.llm.service.ChatConversationMessageService;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import java.time.LocalDateTime;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

@Component
public class CustomerChatMemoryRepository implements ChatMemoryRepository {

    @Autowired
    private ChatConversationMessageService chatConversationMessageService;


    @Override
    public List<String> findConversationIds() {
        return null;
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        return chatConversationMessageService
                .lambdaQuery()
                .eq(ChatConversationMessage::getConversationUuid, conversationId)
                .orderByAsc(ChatConversationMessage::getCreateTime)
                .list()
                .stream().map(this::buildMessage).collect(Collectors.toList());
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void saveAll(String conversationId, List<Message> messages) {
        List<ChatConversationMessage> conversationMessages = messages.stream().map(item -> buildMessage(conversationId, item)).toList();
        deleteByConversationId(conversationId);
        chatConversationMessageService.saveBatch(conversationMessages);
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        chatConversationMessageService
                .lambdaUpdate()
                .eq(ChatConversationMessage::getConversationUuid, conversationId)
                .remove();
    }

    private ChatConversationMessage buildMessage(String conversationId, Message message) {
        LocalDateTime now = LocalDateTime.now();
        String text = message.getText();
        ChatConversationMessage.ChatConversationMessageBuilder builder = ChatConversationMessage.builder()
                .conversationUuid(conversationId)
                .createTime(now)
                .messageType(message.getMessageType());

        if (message.getMessageType().equals(MessageType.ASSISTANT)) {
            String regex = "(<think>.*?</think>)";
            Pattern pattern = Pattern.compile(regex, Pattern.DOTALL);
            Matcher matcher = pattern.matcher(text);
            if (matcher.find()) {
                String thinkMessage = matcher.group(1);
                builder.thinkMessage(thinkMessage);
                builder.message(matcher.replaceFirst(""));
            } else {
                builder.message(text);
            }
        } else {
            builder.message(text);
        }

        return builder.build();

    }



    private Message buildMessage(ChatConversationMessage chatConversationMessage) {
        String message = chatConversationMessage.getMessage();

        MessageType messageType = chatConversationMessage.getMessageType();
        if (messageType.equals(MessageType.USER)) {
            return new UserMessage(message);
        } else if (messageType.equals(MessageType.ASSISTANT)) {
            String thinkMessage = StrUtil.isNotBlank(chatConversationMessage.getThinkMessage()) ? chatConversationMessage.getThinkMessage() : "";
            return new AssistantMessage(thinkMessage + message);
        } else if (messageType.equals(MessageType.SYSTEM)) {
            return new SystemMessage(message);
        } else if (messageType.equals(MessageType.TOOL)) {
            // 现在还没有工具类消息保存
            return null;
//            return new ToolResponseMessage(message)
        } else {
            return null;
        }
    }




}
